import torch.nn as nn
import torch
import numpy as np
import torch.nn.functional as F
from model.decoder import linear_Decoder

    
class TSTEncoder(nn.Module):
    def __init__(self, d_model, n_heads, d_ff=None, 
                        norm='BatchNorm', attn_dropout=0., dropout=0., activation='gelu',
                        res_attention=False, n_layers=1, pre_norm=False, store_attn=False):
        super().__init__()

        self.layers = nn.ModuleList([TSTEncoderLayer(d_model, n_heads=n_heads, d_ff=d_ff, norm=norm,
                                                      attn_dropout=attn_dropout, dropout=dropout,
                                                      activation=activation, res_attention=res_attention,
                                                      pre_norm=pre_norm, store_attn=store_attn) for i in range(n_layers)])
        self.res_attention = res_attention

    def forward(self, src):
        """
        src: tensor [bs x q_len x d_model]
        """
        output = src
        scores = None
        if self.res_attention:
            for mod in self.layers: output, scores = mod(output, prev=scores)
            return output
        else:
            for mod in self.layers: output = mod(output)
            return output

class ScaledDotProductAttention(nn.Module):
    r"""Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer
    (Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets
    by Lee et al, 2021)"""

    def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False):
        super().__init__()
        self.attn_dropout = nn.Dropout(attn_dropout)
        self.res_attention = res_attention
        head_dim = d_model // n_heads
        self.scale = nn.Parameter(torch.tensor(head_dim ** -0.5), requires_grad=lsa)
        self.lsa = lsa

    def forward(self, q, k, v, prev=None, key_padding_mask=None, attn_mask=None):
        '''
        Input shape:
            q               : [bs x n_heads x max_q_len x d_k]
            k               : [bs x n_heads x d_k x seq_len]
            v               : [bs x n_heads x seq_len x d_v]
            prev            : [bs x n_heads x q_len x seq_len]
            key_padding_mask: [bs x seq_len]
            attn_mask       : [1 x seq_len x seq_len]
        Output shape:
            output:  [bs x n_heads x q_len x d_v]
            attn   : [bs x n_heads x q_len x seq_len]
            scores : [bs x n_heads x q_len x seq_len]
        '''

        # Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
        attn_scores = torch.matmul(q, k) * self.scale      # attn_scores : [bs x n_heads x max_q_len x q_len]

        # Add pre-softmax attention scores from the previous layer (optional)
        if prev is not None: attn_scores = attn_scores + prev

        # Attention mask (optional)
        if attn_mask is not None:                                     # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
            if attn_mask.dtype == torch.bool:
                attn_scores.masked_fill_(attn_mask, -np.inf)
            else:
                attn_scores += attn_mask

        # Key padding mask (optional)
        if key_padding_mask is not None:                              # mask with shape [bs x q_len] (only when max_w_len == q_len)
            attn_scores.masked_fill_(key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)

        # normalize the attention weights
        attn_weights = F.softmax(attn_scores, dim=-1)                 # attn_weights   : [bs x n_heads x max_q_len x q_len]
        attn_weights = self.attn_dropout(attn_weights)

        # compute the new values given the attention weights
        output = torch.matmul(attn_weights, v)                        # output: [bs x n_heads x max_q_len x d_v]

        if self.res_attention: return output, attn_weights, attn_scores
        else: return output, attn_weights

class Transpose(nn.Module):
    def __init__(self, *dims, contiguous=False): 
        super().__init__()
        self.dims, self.contiguous = dims, contiguous
    def forward(self, x):        
        if self.contiguous: return x.transpose(*self.dims).contiguous()
        else: return x.transpose(*self.dims)

def positional_encoding(pe, learn_pe, q_len, d_model):
    # Positional encoding
    if pe == None:
        W_pos = torch.empty((q_len, d_model)) # pe = None and learn_pe = False can be used to measure impact of pe
        nn.init.uniform_(W_pos, -0.02, 0.02)
        learn_pe = False
    elif pe == 'zero':
        W_pos = torch.empty((q_len, 1))
        nn.init.uniform_(W_pos, -0.02, 0.02)
    elif pe == 'zeros':
        W_pos = torch.empty((q_len, d_model))
        nn.init.uniform_(W_pos, -0.02, 0.02)
    elif pe == 'normal' or pe == 'gauss':
        W_pos = torch.zeros((q_len, 1))
        torch.nn.init.normal_(W_pos, mean=0.0, std=0.1)
    elif pe == 'uniform':
        W_pos = torch.zeros((q_len, 1))
        nn.init.uniform_(W_pos, a=0.0, b=0.1)
    elif pe == 'sincos': W_pos = PositionalEncoding(q_len, d_model, normalize=True)
    else: raise ValueError(f"{pe} is not a valid pe (positional encoder. Available types: 'gauss'=='normal', \
        'zeros', 'zero', uniform', 'sincos', None.)")
    return nn.Parameter(W_pos, requires_grad=learn_pe)

def PositionalEncoding(q_len, d_model, normalize=True):
    pe = torch.zeros(q_len, d_model)
    position = torch.arange(0, q_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    if normalize:
        pe = pe - pe.mean()
        pe = pe / (pe.std() * 10)
    return pe

SinCosPosEncoding = PositionalEncoding

def get_activation_fn(activation):
    if callable(activation): return activation()
    elif activation.lower() == "relu": return nn.ReLU()
    elif activation.lower() == "gelu": return nn.GELU()
    raise ValueError(f'{activation} is not available. You can use "relu", "gelu", or a callable')

class MultiheadAttention(nn.Module):
    def __init__(self, d_model, n_heads, d_k=None, d_v=None, res_attention=False, attn_dropout=0., proj_dropout=0., qkv_bias=True, lsa=False):
        """Multi Head Attention Layer
        Input shape:
            Q:       [batch_size (bs) x max_q_len x d_model]
            K, V:    [batch_size (bs) x q_len x d_model]
            mask:    [q_len x q_len]
        """
        super().__init__()
        d_k = d_model // n_heads if d_k is None else d_k
        d_v = d_model // n_heads if d_v is None else d_v

        self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v

        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)

        # Scaled Dot-Product Attention (multiple heads)
        self.res_attention = res_attention
        self.sdp_attn = ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout, res_attention=self.res_attention, lsa=lsa)

        # Poject output
        self.to_out = nn.Sequential(nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))


    def forward(self, Q, K=None, V=None, prev=None,
                key_padding_mask=None, attn_mask=None):

        bs = Q.size(0)
        if K is None: K = Q
        if V is None: V = Q

        # Linear (+ split in multiple heads)
        q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,2)       # q_s    : [bs x n_heads x max_q_len x d_k]
        k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0,2,3,1)     # k_s    : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
        v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1,2)       # v_s    : [bs x n_heads x q_len x d_v]

        # Apply Scaled Dot-Product Attention (multiple heads)
        if self.res_attention:
            output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
        else:
            output, attn_weights = self.sdp_attn(q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
        # output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]

        # back to the original inputs dimensions
        output = output.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * self.d_v) # output: [bs x q_len x n_heads * d_v]
        output = self.to_out(output)

        if self.res_attention: return output, attn_weights, attn_scores
        else: return output, attn_weights

class TSTEncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff=256, store_attn=False,
                 norm='BatchNorm', attn_dropout=0, dropout=0., bias=True, 
                activation="gelu", res_attention=False, pre_norm=False):
        super().__init__()
        assert not d_model%n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
        d_k = d_model // n_heads
        d_v = d_model // n_heads

        # Multi-Head attention
        self.res_attention = res_attention
        self.self_attn = MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout, proj_dropout=dropout, res_attention=res_attention)

        # Add & Norm
        self.dropout_attn = nn.Dropout(dropout)
        if "batch" in norm.lower():
            self.norm_attn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
        else:
            self.norm_attn = nn.LayerNorm(d_model)

        # Position-wise Feed-Forward
        self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias),
                                get_activation_fn(activation),
                                nn.Dropout(dropout),
                                nn.Linear(d_ff, d_model, bias=bias))

        # Add & Norm
        self.dropout_ffn = nn.Dropout(dropout)
        if "batch" in norm.lower():
            self.norm_ffn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
        else:
            self.norm_ffn = nn.LayerNorm(d_model)

        self.pre_norm = pre_norm
        self.store_attn = store_attn


    def forward(self, src, prev=None):
        """
        src: tensor [bs x q_len x d_model]
        """
        # Multi-Head attention sublayer
        if self.pre_norm:
            src = self.norm_attn(src)
        ## Multi-Head attention
        if self.res_attention:
            src2, attn, scores = self.self_attn(src, src, src, prev)
        else:
            src2, attn = self.self_attn(src, src, src)
        if self.store_attn:
            self.attn = attn
        ## Add & Norm
        src = src + self.dropout_attn(src2) # Add: residual connection with residual dropout
        if not self.pre_norm:
            src = self.norm_attn(src)

        # Feed-forward sublayer
        if self.pre_norm:
            src = self.norm_ffn(src)
        ## Position-wise Feed-Forward
        src2 = self.ff(src)
        ## Add & Norm
        src = src + self.dropout_ffn(src2) # Add: residual connection with residual dropout
        if not self.pre_norm:
            src = self.norm_ffn(src)

        if self.res_attention:
            return src, scores
        else:
            return src

class PatchTSTEncoder(nn.Module):
    def __init__(self, input_dims, num_patch, patch_len, 
                 n_layers=3, d_model=128, n_heads=16, shared_embedding=True,
                 d_ff=256, n_var=1,norm='BatchNorm', attn_dropout=0., dropout=0., act="gelu",):
        super().__init__()
        self.input_dims = input_dims
        self.num_patch = num_patch
        self.patch_len = patch_len
        self.n_layers = n_layers
        self.hidden_dims = d_model
        self.n_heads = n_heads
        self.n_vars = n_var
        self.share_embedding = shared_embedding
        self.W_pos = positional_encoding("zeros",True,num_patch,d_model)
        self.dropout = nn.Dropout(dropout)
        if not shared_embedding: 
            self.W_P = nn.ModuleList()
            for _ in range(self.n_vars): self.W_P.append(nn.Linear(patch_len, d_model))
        else:
            self.W_P = nn.Linear(patch_len, d_model)     
        self.encoder = TSTEncoder(d_model,n_heads,d_ff,norm,attn_dropout,dropout,activation=act,
                                n_layers = n_layers)


    def forward(self, x):
        x = x.transpose(2,3) # bs * patch * seq * var -> bs * patch * var * seq
        bs,num_patch,n_var,patch_len= x.shape
        if not self.share_embedding:
            x_out = []
            for i in range(n_var):
                z = self.W_P(x[:,:,i,:])
                x_out.append(z)
            x = torch.stack(x_out,dim=2)
        else:
          
            x = self.W_P(x)
        x = x.transpose(1,2)

        u = torch.reshape(x,(bs*n_var,num_patch,self.hidden_dims))
        u = self.dropout(u + self.W_pos)

        z = self.encoder(u)
        z = torch.reshape(z,(-1, n_var, num_patch,self.hidden_dims))
        z = z.permute(0,1,3,2)

        return x
    
class Time_Encoder_Trans(nn.Module):
    def __init__(self,input_dims,output_dims,hidden_dims,batch_size,device,num_patch,patch_len,
                 n_layers,n_heads,shared_embedding,d_ff,n_var):
        super(Time_Encoder_Trans,self).__init__()
        self.input_dims = input_dims
        self.output_dims = output_dims
        self.hidden_dims = hidden_dims
        self.batch_size = batch_size
        self.device = device
    
        self.encoder = PatchTSTEncoder(input_dims,num_patch,patch_len,n_layers,hidden_dims,n_heads,
                                       shared_embedding,d_ff,n_var)

        self.projecty = nn.Sequential(
            nn.Linear(hidden_dims, hidden_dims),
            nn.BatchNorm1d(hidden_dims),
            nn.ReLU(),
            nn.Linear(hidden_dims,output_dims)
        )
    def forward(self,x):
        x = x.to(self.device)
        # print(x.size()) # batch_size * patch * seq_length * feature
        _,patch,_,_ = x.size()
        x = self.encoder(x)        # bs * feature * patch * seq_len
        # x = x.permute(0,1,3,2)  # bs * feature * seq_len * patch
        # print("here",x.squeeze(1).size())
        x = self.projecty(x.view(self.batch_size * patch,self.hidden_dims)).view(self.batch_size,patch,self.output_dims)

        return x



class Time_Encoder(nn.Module):
    def __init__(self,input_dims, output_dims,hidden_dims,batch_size,device):
        super(Time_Encoder,self).__init__()
        self.input_dims = input_dims
        self.output_dims = output_dims
        self.hidden_dims = hidden_dims
        self.batch_size = batch_size
        self.device = device
        self.embedding = nn.Sequential(
            nn.Linear(input_dims,hidden_dims),
            nn.Linear(hidden_dims,hidden_dims),
            nn.Linear(hidden_dims,output_dims)
        )
    def forward(self,x):
        x = x.to(self.device)
        x = self.embedding(x)
        return x

class Fre_Encoder(nn.Module):
    def __init__(self,input_dims, output_dims,hidden_dims,batch_size,device):
        super(Fre_Encoder,self).__init__()
        self.input_dims = input_dims
        self.output_dims = output_dims
        self.hidden_dims = hidden_dims
        self.batch_size = batch_size
        self.device = device
        self.embedding = nn.Sequential(
            nn.Linear(input_dims,hidden_dims),

            nn.Linear(hidden_dims,hidden_dims),

            nn.Linear(hidden_dims,output_dims)
        )
    def forward(self,x):
        x = self.embedding(x)
        return x

class Fre_Encoder_Trans(nn.Module):
    def __init__(self,input_dims,output_dims,hidden_dims,batch_size,device,num_patch,patch_len,
                 n_layers,n_heads,shared_embedding,d_ff,n_var):
        super(Fre_Encoder_Trans,self).__init__()
        self.input_dims = input_dims
        self.output_dims = output_dims
        self.hidden_dims = hidden_dims
        self.batch_size = batch_size
        self.device = device
    
        self.encoder = PatchTSTEncoder(input_dims,num_patch,patch_len,n_layers,hidden_dims,n_heads,
                                       shared_embedding,d_ff,n_var)

        self.project = nn.Sequential(
            nn.Linear(hidden_dims, hidden_dims),
            nn.BatchNorm1d(hidden_dims),
            nn.ReLU(),
            nn.Linear(hidden_dims,output_dims)
        )
    def forward(self,x):
        x = x.to(self.device)
        # print(x.size()) # batch_size * patch * seq_length * feature
        _,patch,_,_ = x.size()
        x = self.encoder(x)
        # print("here",x.size()," ", self.batch_size * patch)
        x = self.project(x.squeeze(1).view(self.batch_size * patch,self.hidden_dims)).view(self.batch_size,patch,self.output_dims)

        return x


class Time_Long_Encoder(nn.Module):
    def __init__(self,input_dims, output_dims,hidden_dims,batch_size,device):
        super(Time_Long_Encoder,self).__init__()
        self.input_dims = input_dims
        self.output_dims = output_dims
        self.hidden_dims = hidden_dims
        self.batch_size = batch_size
        self.device = device
        self.embedding = nn.Sequential(
            nn.Linear(input_dims,hidden_dims),

            nn.Linear(hidden_dims,hidden_dims),
            nn.Linear(hidden_dims,output_dims)
        )
    def forward(self,x):
        x = self.embedding(x)
        return x

class Time_Long_Encoder_Trans(nn.Module):
    def __init__(self,input_dims, output_dims,hidden_dims,seq_len,batch_size,device):
        super(Time_Long_Encoder_Trans,self).__init__()
        self.input_dims = input_dims
        self.output_dims = output_dims
        self.hidden_dims = hidden_dims
        self.batch_size = batch_size
        self.seq_len = seq_len
        self.device = device
        self.encoder_layer = nn.TransformerEncoderLayer(self.hidden_dims,8,output_dims)
        self.embedding = nn.TransformerEncoder(self.encoder_layer,2)
        self.project = nn.Sequential(
            nn.Linear(self.seq_len,self.hidden_dims),
            nn.Linear(self.hidden_dims,self.hidden_dims)
        )

    def forward(self,x):
        # print(x.size())
        x = self.project(x.view(self.batch_size,self.seq_len))
        x = self.embedding(x)
        # print("x.shape : ",x.size())
        return x


class Time_Frequence_Mul(nn.Module):
    def __init__(self,input_dims,output_dims,hidden_dims,lr,batch_size,device,var,configs):
        super(Time_Frequence_Mul,self).__init__()
        self.input_dims = input_dims
        self.output_dims = output_dims
        self.hidden_dims = hidden_dims
        self.device = device
        self.lr = lr
        self.batch_size = batch_size
        self.configs = configs
        # path_len,n_layers,n_heads,shared_embedding,d_ff
        # self.time_encoder = Time_Encoder(self.input_dims,self.output_dims,self.hidden_dims,self.batch_size,self.device)
        self.time_encoder = Time_Encoder_Trans(self.input_dims,self.output_dims,self.hidden_dims,self.batch_size,self.device,
                            self.configs.PATCHES,self.configs.PATCH_LEN,self.configs.ENCODER_LAYERS,self.configs.ENCODER_HEAD,
                            self.configs.SHARED_EMBDEDDING,self.configs.D_FF,self.configs.VARS)
        self.fre_encoder = Fre_Encoder_Trans(self.input_dims,self.output_dims,self.hidden_dims,self.batch_size,self.device,
                            self.configs.PATCHES,self.configs.PATCH_LEN,self.configs.ENCODER_LAYERS,self.configs.ENCODER_HEAD,
                            self.configs.SHARED_EMBDEDDING,self.configs.D_FF,self.configs.VARS)
        self.time_long_encoder = Time_Long_Encoder_Trans(self.input_dims,self.output_dims,self.hidden_dims,self.configs.SEQ_LEN,
                                                         self.batch_size,self.device)
        self.decoder = linear_Decoder(configs.SEQ_LEN,configs.LABEL_LEN + configs.PRED_LEN,configs.HIDDEN_DIMS,configs.lr,configs.BATCH_SIZE,configs.DEVICE,configs.VARS)
        
    def fit(self, x, feq_x, long_x):
        # 每个Encoder都对
        time_x = self.time_encoder(x)
        feq_x = self.fre_encoder(feq_x)
        time_long_x = self.time_long_encoder(long_x)
        pred = self.decoder(time_long_x)
        return [time_x,feq_x,time_long_x,pred]


